'''
Author: goog
Date: 2021-12-15 12:31:41
LastEditTime: 2022-01-14 09:49:01
LastEditors: goog
Description: 定位ls
FilePath: /chengdu/TensorRT20220110/DetectionFM/RD/RD.py
Time Limit Exceeded!
'''

import cv2
import json
import os
import numpy as np
import torch
from torchvision import transforms
from PIL import Image

def result_cls(cfg, net, roi):
    device = torch.device(cfg['device'])
    imgtransform = transforms.Compose([
            transforms.ToTensor(),
            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
    ])
    roi = cv2.resize(roi, (100, 100)) # 这个是关键 使用opencv的resize 使用PIL的导致错误
    img = Image.fromarray(roi)
    img = imgtransform(img)
    img = img.unsqueeze(0)
    img = img.to(device)
    with torch.no_grad():
        result = net(img)  # 1*3
        result = torch.softmax(result, dim=1)  
        # print(result)          
        state = int(result.argmax())
        # output = torch.squeeze(net(img)).cpu()
        # # # print(output.cpu().data.numpy())
        # predict = torch.softmax(output, dim=0)
        # print(predict)
        # predict_cla = torch.argmax(predict).numpy()
        # state = predict_cla
        return state


def rd(cfg, image1, item1, barcode, cam_ip, det_net, visualization=False):
    image = image1.copy()
    image_roi = cv2.cvtColor(image1.copy(), cv2.COLOR_BGR2RGB)
    defect = "RD"
    

    if 'RR-L' in barcode:
        bias8 = -300
        bias7 = -550
    elif 'RR-R' in barcode:
        bias8 = -60
        bias7 = -220
    elif 'FR-L' in barcode:
        bias8 = -20
        bias7 = -220
    elif 'FR-R' in barcode:
        bias8 = -20
        bias7 = -250
    else:
        bias8 = -20
        bias7 = -220
        print('rd error')
    

    abs_dir = cfg['abs_dir']


    # load config
    config_path = os.path.join(abs_dir, 'DetectionFM/'+defect+'/'+defect+'.json')
    with open(config_path, 'r') as fp:
        config = json.load(fp)
        if barcode not in config.keys():
            barcode = 'C50'+barcode[3:]
    defect_count = config[barcode][cam_ip] 
    
    item_tp = item1['tp']
    if len(item_tp)>0:

        item_defect = item1[defect.lower()]
        defect_np = np.array(item_defect)
        if cam_ip == "192.168.8.8":
            plot_bias = bias8
            ind = np.where((defect_np[:, 2]<item_tp[0][2]+bias8)&(defect_np[:, 2]>1000))
        else:
            ind = np.where((defect_np[:, 2]>(item_tp[0][2]+bias7))&(defect_np[:, 2]<4000))
            plot_bias = bias7
        defect_np_lt = defect_np[ind[0], :]


        defect_np_lt = defect_np[ind[0], :]

        # conf setting
        conf = config[barcode]['conf']
        conf_ind = np.where(defect_np_lt[:, -1]>conf)   
        # print(conf_ind) 
        defect_np_lt = defect_np_lt[conf_ind[0], :]

        # visual
        object_count = 0
        defect_object_count = 0
        if visualization:
            for v in defect_np_lt:
                img = cv2.line(image, 
                            (item_tp[0][2]+plot_bias, 0), 
                            (item_tp[0][2]+plot_bias, 3647), 
                            color=(255, 0 ,255),
                            thickness=10)
                xmin, ymin, xmax, ymax = int(v[0]), int(v[1]), int(v[2]), int(v[3])
                state = result_cls(cfg, net=det_net, roi=image_roi[ymin-10:ymax+10, xmin-10:xmax+10, :])
                if state:
                    img = cv2.rectangle(img, 
                                        (xmin, ymin), 
                                        (xmax, ymax), 
                                        color=(0, 255, 0),
                                    thickness=6)
                    object_count += 1
                else:
                    defect_object_count +=1
                    img = cv2.rectangle(img, 
                                        (xmin, ymin), 
                                        (xmax, ymax), 
                                        color=(0, 0, 255),
                                    thickness=12)
            cv2.imwrite(os.path.join(abs_dir, 'Test', defect+cam_ip+'.jpg'), img=img)
            print("{}_{}\tyolo:{}\\{}\t\tcls:noraml[{}]\tabnormal[{}]:".format(cam_ip, defect, defect_np_lt.shape[0], defect_count, object_count, defect_object_count))

        
               

        # return defect_np_lt.shape[0]==defect_count, img
        if 'FR-L' in barcode and cam_ip == "192.168.8.8":
            return object_count == defect_count-2, img
        elif 'FR-L' in barcode and cam_ip == "192.168.7.7":
            return object_count == defect_count-1, img
        elif 'FR-R' in barcode and cam_ip == "192.168.8.8":
            return object_count == defect_count-2, img
        elif 'FR-R' in barcode and cam_ip == "192.168.7.7":
            return object_count == defect_count-1, img
        else:
            return object_count == defect_count, img

